from data_utils import to_batches
from tqdm import tqdm
import re
import pandas as pd
import json
import requests
import numpy as np


def get_zpoint_embedding(texts: list[str]) -> np.ndarray:
    res = requests.post(url="http://172.16.3.123:7777/embeddings/", data=json.dumps(texts))
    embeddings = json.loads(res.content.decode("utf-8"))["embeddings"]
    result = np.array(embeddings, dtype='float32')
    return result


def split_table_from_content(rc_splitter, content):
    split_ls = []
    for e in re.split('(<table:[^>]+>)', content):
        table_mds = re.findall('<table:([^>]+)>', e)
        if table_mds:
            split_ls.extend(table_mds)
        else:
            if e != '\n' and len(e.strip()) > 0:
                split_ls.extend(rc_splitter.split_text(e))
    return split_ls


def split_docx_content(rc_splitter, path):
    with open(path, 'r', encoding='utf-8') as f:
        content_ls = json.load(f)
    df_ls = []
    for dic in content_ls:
        doc_content = dic['content']
        img_info = '\n'.join(dic['img_path'])
        file_uuid = dic['file_uuid']
        split_blocks = split_table_from_content(rc_splitter, doc_content)
        df_dict = {'file_uuid': [file_uuid] * len(split_blocks), 'text': split_blocks,
                   'image_path': [img_info] * len(split_blocks)}
        df_ls.append(pd.DataFrame(df_dict, index=None))
    df = pd.concat(df_ls, axis=0)
    return df


def add_embeddings(df):
    df_dict = df.to_dict(orient='records')
    texts = df['text'].tolist()
    print('max_len: ', max([len(txt) for txt in texts]))
    batch_size = 20
    text_batches = to_batches(texts, batch_size)
    text_embeddings = []
    bar1 = tqdm(text_batches)
    for batch in bar1:
        embeddings = get_zpoint_embedding(batch)
        text_embeddings.extend(embeddings)
        bar1.set_description('计算正文Embeddings')
    file_uuids = df['file_uuid'].tolist()
    file_uuid_batches = to_batches(file_uuids, batch_size)
    file_uuid_embeddings = []
    bar2 = tqdm(file_uuid_batches)
    for batch in bar2:
        embeddings = get_zpoint_embedding(batch)
        file_uuid_embeddings.extend(embeddings)
        bar2.set_description('计算文件名Embeddings')
    for i, item in enumerate(df_dict):
        item['text_embedding'] = text_embeddings[i]
        item['file_uuid_embedding'] = file_uuid_embeddings[i]
    return df_dict
